import os
from datetime import datetime

import matplotlib.pyplot as plt


class TrainingMonitor:
    def __init__(self, model_name):
        self.model_name = model_name
        self.train_losses = []
        self.val_losses = []
        self.train_accs = []
        self.val_accs = []
        self.timestamps = []

    def record(self, epoch, train_loss, val_loss, train_acc, val_acc):
        self.train_losses.append(train_loss)
        self.val_losses.append(val_loss)
        self.train_accs.append(train_acc)
        self.val_accs.append(val_acc)
        self.timestamps.append(datetime.now().strftime("%H:%M:%S"))

    def plot_curves(self, save_path='training_curves.png'):
        plt.figure(figsize=(12, 6))

        # 损失曲线
        plt.subplot(1, 2, 1)
        plt.plot(self.train_losses, 'r-o', label='Train Loss')
        plt.plot(self.val_losses, 'b-s', label='Val Loss')
        plt.title(f'{self.model_name} Loss Curve')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.legend()

        # 准确率曲线
        plt.subplot(1, 2, 2)
        plt.plot(self.train_accs, 'g-o', label='Train Acc')
        plt.plot(self.val_accs, 'm-s', label='Val Acc')
        plt.title(f'{self.model_name} Accuracy Curve')
        plt.xlabel('Epoch')
        plt.ylabel('Accuracy')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.legend()

        # 保存和显示
        plt.tight_layout()
        os.makedirs('training_plots', exist_ok=True)
        plt.savefig(f'training_plots/{self.model_name}_{save_path}')
        plt.close()
